"""
Run Burgers equation experiments from Section 3 of the paper.
"""

import sys
import os
import numpy as np
import matplotlib.pyplot as plt

# Add parent directory to path
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..'))

from locELM.examples.burgers import BurgersEquationProblem


def run_base_experiment():
    """
    Run the base experiment: 40 time blocks, t_f=10
    Figure ref: fig:burger
    """
    print("\n" + "="*70)
    print("EXPERIMENT 1: Base Burgers Equation Test (t_f=10, 40 time blocks)")
    print("="*70)

    problem = BurgersEquationProblem(
        a=0.0,
        b=5.0,
        t_f=10.0,
        nu=0.01,
        n_time_blocks=40,
        n_x_subdomains=5,
        n_t_subdomains=1,
        q_x=20,
        q_t=20,
        n_hidden=200,
        rand_mag=0.75,
        random_seed=22
    )

    results = problem.solve()

    # Save results
    output_dir = os.path.join(os.path.dirname(__file__), '..', 'experiments')
    os.makedirs(output_dir, exist_ok=True)

    np.savez(
        os.path.join(output_dir, 'burgers_base_experiment.npz'),
        overall_max_error=results['overall_max_error'],
        overall_rms_error=results['overall_rms_error'],
        total_time=results['total_time'],
        n_time_blocks=40,
        n_x_subdomains=5,
        q_x=20,
        q_t=20,
        n_hidden=200,
        rand_mag=0.75
    )

    return results


def run_dof_variation_experiment():
    """
    Run experiments varying degrees of freedom.
    Figure ref: fg_bg_3
    t_f=2.5, 10 time blocks
    """
    print("\n" + "="*70)
    print("EXPERIMENT 2: DOF Variation Study (t_f=2.5, 10 time blocks)")
    print("="*70)

    # Test different numbers of collocation points (fixed M=200)
    print("\n--- Part A: Varying collocation points (fixed M=200) ---")
    colloc_points = [5, 10, 15, 20, 25]
    colloc_results = []

    for q in colloc_points:
        print(f"\nTesting Q={q}x{q} collocation points...")
        problem = BurgersEquationProblem(
            a=0.0,
            b=5.0,
            t_f=2.5,
            nu=0.01,
            n_time_blocks=10,
            n_x_subdomains=5,
            n_t_subdomains=1,
            q_x=q,
            q_t=q,
            n_hidden=200,
            rand_mag=0.5,
            random_seed=22
        )
        result = problem.solve()
        colloc_results.append({
            'q': q,
            'max_error': result['overall_max_error'],
            'rms_error': result['overall_rms_error'],
            'time': result['total_time']
        })

    # Test different numbers of training parameters (fixed Q=20x20)
    print("\n--- Part B: Varying training parameters (fixed Q=20x20) ---")
    training_params = [50, 100, 150, 200, 250, 300]
    training_results = []

    for m in training_params:
        print(f"\nTesting M={m} training parameters...")
        problem = BurgersEquationProblem(
            a=0.0,
            b=5.0,
            t_f=2.5,
            nu=0.01,
            n_time_blocks=10,
            n_x_subdomains=5,
            n_t_subdomains=1,
            q_x=20,
            q_t=20,
            n_hidden=m,
            rand_mag=0.5,
            random_seed=22
        )
        result = problem.solve()
        training_results.append({
            'm': m,
            'max_error': result['overall_max_error'],
            'rms_error': result['overall_rms_error'],
            'time': result['total_time']
        })

    # Save results
    output_dir = os.path.join(os.path.dirname(__file__), '..', 'experiments')
    np.savez(
        os.path.join(output_dir, 'burgers_dof_variation.npz'),
        colloc_results=colloc_results,
        training_results=training_results
    )

    return colloc_results, training_results


def run_subdomain_variation_experiment():
    """
    Run experiments varying number of subdomains with fixed total DOF.
    Figure ref: fg_bg_4
    t_f=0.25, 1 time block
    """
    print("\n" + "="*70)
    print("EXPERIMENT 3: Subdomain Variation (fixed total DOF, t_f=0.25)")
    print("="*70)

    # Different configurations with ~2000 total collocation points
    # and ~1000 total training parameters
    configs = [
        {'n_x': 1, 'q': 45, 'm': 1000, 'r_m': 2.0},
        {'n_x': 2, 'q': 32, 'm': 500, 'r_m': 1.0},
        {'n_x': 3, 'q': 26, 'm': 333, 'r_m': 1.0},
        {'n_x': 4, 'q': 22, 'm': 250, 'r_m': 0.75},
        {'n_x': 5, 'q': 20, 'm': 200, 'r_m': 0.75},
    ]

    subdomain_results = []

    for config in configs:
        n_x = config['n_x']
        q = config['q']
        m = config['m']
        r_m = config['r_m']

        print(f"\nTesting N_x={n_x}, Q={q}x{q}, M={m}, R_m={r_m}...")
        problem = BurgersEquationProblem(
            a=0.0,
            b=5.0,
            t_f=0.25,
            nu=0.01,
            n_time_blocks=1,
            n_x_subdomains=n_x,
            n_t_subdomains=1,
            q_x=q,
            q_t=q,
            n_hidden=m,
            rand_mag=r_m,
            random_seed=22
        )
        result = problem.solve()
        subdomain_results.append({
            'n_x': n_x,
            'q': q,
            'm': m,
            'r_m': r_m,
            'max_error': result['overall_max_error'],
            'rms_error': result['overall_rms_error'],
            'time': result['total_time']
        })

    # Save results
    output_dir = os.path.join(os.path.dirname(__file__), '..', 'experiments')
    np.savez(
        os.path.join(output_dir, 'burgers_subdomain_variation.npz'),
        subdomain_results=subdomain_results
    )

    return subdomain_results


def main():
    """Run all Burgers equation experiments."""
    # Run base experiment
    base_results = run_base_experiment()

    # Run DOF variation experiments
    colloc_results, training_results = run_dof_variation_experiment()

    # Run subdomain variation experiment
    subdomain_results = run_subdomain_variation_experiment()

    # Print summary
    print("\n" + "="*70)
    print("ALL EXPERIMENTS COMPLETED")
    print("="*70)
    print(f"\nBase experiment max error: {base_results['overall_max_error']:.3e}")
    print(f"Results saved to experiments/ directory")


if __name__ == '__main__':
    main()
